package cn.iocoder.boot.MLlib;

import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.streaming.api.java.JavaReceiverInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.regression.LinearRegressionModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import scala.Tuple2;

import java.util.ArrayList;
import java.util.List;

@Service
//Spring 组件类，用于处理实时数据流并集成机器学习。
public class DataProcessingService {

    @Autowired
    private JavaStreamingContext javaStreamingContext;

    @Autowired
    private SparkSession sparkSession;

    public void processDataStream() {
        // 从数据源（如 Socket）接收数据，得到 JavaReceiverInputDStream<String>
        JavaReceiverInputDStream<String> lines = javaStreamingContext.socketTextStream("localhost", 9999);

        // 使用 foreachRDD 来处理每个 batch 的数据
        lines.foreachRDD(rdd -> {
            // 将 JavaReceiverInputDStream 的每个 RDD 转换为 JavaPairRDD<String, String>
            JavaPairRDD<String, String> pairRDD = rdd.mapToPair(line -> {
                // 假设每行文本数据包含一个键和一个值，用逗号分隔
                String[] parts = line.split(",");
                return new Tuple2<>(parts[0], parts[1]);
            });

            // 对数据进行预处理，如清洗、转换等
            JavaPairRDD<String, Double> processedData = pairRDD.mapValues(value -> {
                // 这里进行数据清洗和转换的具体逻辑
                return Double.parseDouble(value);
            });

            // 将处理后的数据转换为 Spark MLlib 的输入格式
            JavaRDD<Row> mlData = processedData.map(tuple -> RowFactory.create(tuple._1(), tuple._2()));

            // 定义特征列和标签列
            List<String> featureColumns = new ArrayList<>();
            featureColumns.add("feature1");
            featureColumns.add("feature2");

            // 创建 VectorAssembler 将特征列组合成一个特征向量
            VectorAssembler assembler = new VectorAssembler()
                    .setInputCols(featureColumns.toArray(new String[0]))
                    .setOutputCol("features");

            // 将数据转换为 DataFrame
            Dataset<Row> dataFrame = sparkSession.createDataFrame(mlData, YourDataClass.class);

            // 应用 VectorAssembler 转换数据
            Dataset<Row> assembledData = assembler.transform(dataFrame);

            // 加载预训练的模型（例如线性回归模型）
            LinearRegressionModel model = LinearRegressionModel.load("path/to/trained/model");

            // 使用模型进行预测
            Dataset<Row> predictions = model.transform(assembledData);

            // 输出预测结果
            predictions.show();
        });

        // 启动流处理
        javaStreamingContext.start();
        try {
            javaStreamingContext.awaitTermination();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}